# encoding: utf-8
#!/usr/bin/python3

import warnings
warnings.filterwarnings("ignore")

import torch
import os
from loss import Heatmap_Classify_Loss
import pandas as pd
from torch.utils.data import DataLoader
import torch.nn as nn
from process_data import conver_csv
from dataset import Mydataset
from val import val_fun
from torch.optim.lr_scheduler import StepLR,CosineAnnealingLR
from warmup_scheduler import GradualWarmupScheduler
from models.model import model_cpn50,model_simple_pose_res18,model_simple_pose_res101,\
    model_unet,model_resnet_unet,model_DLinknet34,model_simple_pose_res34
from data_aug import tensor_resize
import time
Model_Save_Patch = '../model_save/new/Dlink34_MISH_768'
Record_File_Dir = './record/new/Dlink34_MISH_768.txt'
Image_Patch = '../data/Train'
Sigma = 9
Image_Size = (768,768)
Batch_Size = 4
lr = 1e-5
num_epoch = 500
accumlation_step = 8
Num_Class = 2
train_csv_dir = '../data/csv_label/train150_anno.csv'
val_csv_dir= '../data/csv_label/train51_anno.csv'
os.makedirs(Model_Save_Patch,exist_ok=True)
os.makedirs(Record_File_Dir.replace(Record_File_Dir.split('/')[-1],''),exist_ok=True)
loss_fun = nn.MSELoss()


def train(model):
    best_result = 100
    if not os.path.exists(Model_Save_Patch):
        os.makedirs(Model_Save_Patch)
    data_list_1 = conver_csv(train_csv_dir)
    data_list_2 = conver_csv(val_csv_dir)
    data_list_train = data_list_1[:130] + data_list_2[:20]
    data_list_val = data_list_1[130:] + data_list_2[20:]

    My_dataSet = Mydataset(anno_list=data_list_train,size = Image_Size ,image_path='../data/image',sigma=Sigma,)
    train_data_loader = DataLoader(dataset=My_dataSet, batch_size=Batch_Size, shuffle=True,
                                   num_workers=4, pin_memory=True)

    optimizer = torch.optim.Adam(model.parameters(),lr = lr, weight_decay= 2e-5)
    # optimizer = torch.optim.SGD(model.parameters(),lr = lr,momentum=0.9,weight_decay=2e-5)
    scheduler = CosineAnnealingLR(optimizer,T_max=num_epoch,)
    scheduler = GradualWarmupScheduler(optimizer, multiplier=10, total_epoch=20, after_scheduler=scheduler)
    total_step = len(train_data_loader)

    # pretrained_dict = torch.load('../model_save/DlinkNet34/best_.pth')
    # model_dict = model.state_dict()
    # pretrained_dict = {k:v for k,v in pretrained_dict.items() if k in model_dict.keys()}
    # # print(pretrained_dict.keys())
    # model_dict.update(pretrained_dict)
    # model.load_state_dict(model_dict)

    for epoch in range(num_epoch):
        # model.load_state_dict(torch.load('../model_save/heatmap/simple_pose_res101/63.pth'))
        model.train()
        start_time  = time.time()
        #注意必须每次训练都设置model.train()，因为在val时会设置model.eval(),
        #需要再次进行转换。

        print('epoch{}'.format(epoch+1))
        scheduler.step()
        train_loss = 0
        for step,(batch_image,batch_label) in enumerate(train_data_loader):
            # pass
            if torch.cuda.is_available():
                batch_image = batch_image.float().cuda()
                batch_label = batch_label.float().cuda()
            # ### 使用pytorch中的函数进行随机缩放
            # batch_image, batch_label = tensor_resize(batch_image,batch_label,train_size=Image_Size)
            # print('batch_image',batch_image.shape,'batch_label',batch_label.shape)
            pred = model(batch_image)
            loss = loss_fun(pred,batch_label)
            # loss = loss_fun(pred[0][0],batch_label)+loss_fun(pred[0][1],batch_label)+\
            #        loss_fun(pred[0][2],batch_label)+loss_fun(pred[0][3],batch_label)+4*loss_fun(pred[1],batch_label)
            loss.backward()
            if (step % accumlation_step) == 0:
                optimizer.step()
                optimizer.zero_grad()
            if step % 20 ==0 and epoch <3:
                print('Step [{}/{}],Loss:{:.7f}'.format(step + 1, total_step, loss))
            train_loss += loss.item()

        train_loss = train_loss/total_step
        print('Epoch [{}/{}],Loss:{:.7f}'.format(epoch+1,num_epoch,train_loss))

        # checkpoint = {
        #     'epoch':epoch+1,
        #     'model':model.state_dict(),
        #     'optimizer':optimizer.state_dict(),
        # }
        if ((epoch % 2 == 0) and (epoch > 60 or epoch < 15)) or (epoch % 4 ==0):
            torch.cuda.empty_cache()
            record_file = open(Record_File_Dir,'a')
            val_recall, val_l2_mean, val_l2_medium, = val_fun(model, data_list_val, Image_Size, distance=1, show_pred=False,)
            train_recall, train_l2_mean, train_l2_medium = val_fun(model, data_list_train, Image_Size, distance=1, show_pred=False,)
            print('val result--epoch',str(epoch+1),'--val_recall',val_recall,'--val_l2_mean--',val_l2_mean,
                  '--val_l2_medium--',val_l2_medium,)
            print('train result--epoch', str(epoch + 1), '--train_recall', train_recall, '--train_l2_mean--', train_l2_mean,
                  '--train_l2_medium--', train_l2_medium)


            record_file.write(
                'epoch--' + str(epoch + 1) + '--train_loss--' + str(train_loss) + '--val_recall--' + str(
                    val_recall) +
                '--val_l2_mean--' + str(val_l2_mean) + '--val_l2_medium--' + str(val_l2_medium))
            record_file.write('\n')
            record_file.write('--train_recall--' + str(train_recall) + '--train_l2_mean--' + str(train_l2_mean)
                              + '--train_l2_medium--' + str(train_l2_medium))
            record_file.write('\n')
            record_file.close()

            torch.cuda.empty_cache()

            torch.save(model.state_dict(), (Model_Save_Patch + '/' + 'new_' + '.pth'))
            if val_l2_medium < best_result:
                print('get better model with medium of ',val_l2_medium)
                torch.save(model.state_dict(), (Model_Save_Patch + '/' + 'best_' + '.pth'))
                best_result = val_l2_medium
            print('---------------------------------------------')


        if epoch < 3:
            print('use time -------',time.time()-start_time)
        torch.save(model.state_dict(),Model_Save_Patch+'/'+str(num_epoch+1)+'.pth')
        torch.cuda.empty_cache()






if __name__ == '__main__':
    torch.backends.cudnn.benchmark = True
    # train(model = model_cpn50(Image_Size),)
    # train(model=model_simple_pose_res18(pretrain=True))
    # train(model_unet())
    # train(model_resnet_unet(layer='resnet34',pre_train=False,num_class=Num_Class))
    # train(model_DLinknet34(num_class=Num_Class,pre_trained=False))
    # train(model_simple_pose_res18(num_class=Num_Class,pretrain=False))
    # train(model_simple_pose_res34(num_class=Num_Class))
    train(model_DLinknet34(num_class=2,pre_trained=False))